package com.example.tree.heap;

import java.util.*;

/**
 * 给一非空的单词列表，返回前 k 个出现次数最多的单词。
 * 返回的答案应该按单词出现频率由高到低排序。如果不同的单词有相同出现频率，按字母顺序排序。
 * <p>
 * 示例 1：
 * <p>
 * 输入: ["i", "love", "leetcode", "i", "love", "coding"], k = 2
 * 输出: ["i", "love"]
 * 解析: "i" 和 "love" 为出现次数最多的两个单词，均为2次。
 * 注意，按字母顺序 "i" 在 "love" 之前。
 * 示例 2：
 * <p>
 * 输入: ["the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"], k= 4
 * 输出: ["the", "is", "sunny", "day"]
 * 解析: "the", "is", "sunny" 和 "day" 是出现次数最多的四个单词，
 * 出现次数依次为 4, 3, 2 和 1 次。
 * <p>
 * <p>
 * <p>
 * 注意：
 * 假定 k 总为有效值， 1 ≤ k ≤ 集合元素数。
 * 输入的单词均由小写字母组成。
 */
public class Leetcode692_TopKFrequent {
    public static void main(String[] args) {
        String[] words = {"the", "day", "is", "sunny", "the", "the", "the", "sunny", "is", "is"};
        int k = 4;

        for (String s : topKFrequent3(words, k)) {
            System.out.println(s);
        }

    }

    /**
     * 最优解:桶排序+Trie树
     * <p>
     * 1.统计每个单词的词频
     * 2.将词频相同的单词放入同一个桶中，每个桶用Trie树存储字符串
     *
     * @param words
     * @param k
     * @return
     */
    public static List<String> topKFrequent3(String[] words, int k) {
        List<String> result = new ArrayList<>();

        // 统计词频
        Map<String, Integer> map = new HashMap<>();
        for (String word : words) {
            map.put(word, map.getOrDefault(word, 0) + 1);
        }

        Trie[] buckets = new Trie[words.length];
        // 遍历每个单词，将其加入它所在桶的trie树
        for (Map.Entry<String, Integer> entry : map.entrySet()) {
            String word = entry.getKey();
            Integer frequent = entry.getValue();
            if (buckets[frequent] == null) {
                buckets[frequent] = new Trie();
            }
            buckets[frequent].addWord(word);
        }

        // 对于每个桶中的trie树，将单词按照字典序排序，和k比较
        for (int i = buckets.length - 1; i >= 0; i--) {

            if (buckets[i] != null) {
                List<String> tmp = buckets[i].getWords();
                if (tmp.size() < k) {
                    result.addAll(tmp);
                    k = k - tmp.size();
                } else {
                    for (int j = 0; j < k; j++) {
                        result.add(tmp.get(j));
                    }
                    break;
                }
            }

        }

        return result;
    }


    private static class Trie {
        private TrieNode root;

        public Trie() {
            root = new TrieNode();
        }

        public void addWord(String word) {
            TrieNode cur = root;
            for (int i = 0; i < word.length(); i++) {
                int pos = word.charAt(i) - 'a';
                if (cur.children[pos] == null) {
                    cur.children[pos] = new TrieNode();
                }
                cur = cur.children[pos];
            }
            cur.word = word;

        }

        public List<String> getWords() {
            List<String> res = new ArrayList<>();
            getWords(this.root, res);
            return res;
        }

        public void getWords(TrieNode root, List<String> res) {
            if (root == null) {
                return;
            }
            if (root.word != null) {
                res.add(root.word);
            }

            for (int i = 0; i < 26; i++) {
                if (root.children[i] != null) {
                    getWords(root.children[i], res);
                }
            }
        }

        class TrieNode {
            TrieNode[] children;
            String word;

            public TrieNode() {
                this.children = new TrieNode[26];
                this.word = null;
            }
        }
    }

    /**
     * 解法二:最小堆
     *
     * @param words
     * @param k
     * @return
     */
    public static List<String> topKFrequent2(String[] words, int k) {
        List<String> result = new ArrayList<>();
        // 统计词频
        Map<String, Integer> map = new HashMap<>();
        for (String word : words) {
            map.put(word, map.getOrDefault(word, 0) + 1);
        }

        // 构建一个容量为K小根堆
        Queue<String> minHeap = new PriorityQueue<>(k, (word1, word2) -> {
            return map.get(word1).equals(map.get(word2)) ?
                    word2.compareTo(word1) : map.get(word1) - map.get(word2);
        });

        for (String s : map.keySet()) {
            minHeap.offer(s);
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }

        while (!minHeap.isEmpty()) {
            result.add(minHeap.poll());
        }

        Collections.reverse(result);

        return result;

    }

    /**
     * 解法一:暴力法(Map + 自定义排序)
     * <p>
     * 使用提供的单词数组的单词为key，单词出现的次数为value形成Map
     * 然后按词频进行排序，返回词频前的单词
     *
     * @param words
     * @param k
     * @return
     */
    public static List<String> topKFrequent1(String[] words, int k) {
        Map<String, Integer> map = new HashMap<>();

        // 统计每个出现单词的词频
        for (String word : words) {
            map.put(word, map.getOrDefault(word, 0) + 1);

        }
        List<String> result = new ArrayList<>(map.keySet());

        // 根据词频进行排序
        Collections.sort(result, (word1, word2) -> {
            return map.get(word1).equals(map.get(word2)) ?
                    word1.compareTo(word2) : map.get(word2) - map.get(word1);
        });

        return result.subList(0, k);
    }
}
